from core.layer import Layer
from implement.layers.basic.affine import Affine
from utils.functions_collect import tanh, sigmoid

'''
这个类实现了长短时记忆网络 (LSTM) 层，其中包含了一系列的仿射变换和激活函数，用于对输入数据进行处理并更新隐藏状态和细胞状态。
'''
class LSTM(Layer):
    def __init__(self, hidden_size, in_size=None):
        """长短时记忆网络 (LSTM) 层。

        Args:
            hidden_size (int): LSTM 中隐藏层的大小。
            in_size (int or None): 输入数据的大小。如果为 None，则在第一次前向传播时根据输入形状确定。
        """
        super().__init__()

        self.c = None
        self.h = None
        H, I = hidden_size, in_size
        self.x2f = Affine(H, in_size=I)
        self.x2i = Affine(H, in_size=I)
        self.x2o = Affine(H, in_size=I)
        self.x2u = Affine(H, in_size=I)
        self.h2f = Affine(H, in_size=H, nobias=True)
        self.h2i = Affine(H, in_size=H, nobias=True)
        self.h2o = Affine(H, in_size=H, nobias=True)
        self.h2u = Affine(H, in_size=H, nobias=True)
        self.reset_state()

    def reset_state(self):
        """重置 LSTM 的状态，即隐藏层和细胞状态。"""
        self.h = None
        self.c = None

    def forward(self, x):
        """LSTM 的前向传播。

        Args:
            x (ndarray): 输入数据。

        Returns:
            h_new (ndarray): LSTM 的输出，即新的隐藏状态。
        """
        if self.h is None:
            f = sigmoid(self.x2f(x))
            i = sigmoid(self.x2i(x))
            o = sigmoid(self.x2o(x))
            u = tanh(self.x2u(x))
        else:
            f = sigmoid(self.x2f(x) + self.h2f(self.h))
            i = sigmoid(self.x2i(x) + self.h2i(self.h))
            o = sigmoid(self.x2o(x) + self.h2o(self.h))
            u = tanh(self.x2u(x) + self.h2u(self.h))

        if self.c is None:
            c_new = (i * u)
        else:
            c_new = (f * self.c) + (i * u)

        h_new = o * tanh(c_new)

        self.h, self.c = h_new, c_new
        return h_new

